#include "ShaderResourcesManager.h"

#include <cassert>
#include <fstream>
#include <D3D11.h>
#include <vector>

#include "main/D3DData.h"
#include "main/Globals.h"
#include "utils/DDSTextureLoader.h"
#include "utils/DxErrorChecker.h"

namespace
{
    void createTexture2DArraySRV(const std::vector<std::wstring>& filenames, ID3D11ShaderResourceView* resourceView) {
        //
        // Load the texture elements individually from file.  These textures
        // won't be used by the GPU (0 bind flags), they are just used to 
        // load the image data from file.  We use the STAGING usage so the
        // CPU can read the resource.
        //

        assert(Globals::gD3D->mDevice);
        ID3D11Device& device = *Globals::gD3D->mDevice;

		assert(Globals::gD3D->mImmediateContext);
		ID3D11DeviceContext& context = *Globals::gD3D->mImmediateContext;

        const uint32_t numFilenames = static_cast<uint32_t> (filenames.size());

        std::vector<ID3D11Texture2D*> sourceTextures;
        sourceTextures.resize(numFilenames);
        for(uint32_t filenameIndex = 0; filenameIndex < numFilenames; ++filenameIndex) {
            ID3D11Texture2D* &resource = sourceTextures[filenameIndex];
            ID3D11Resource* &texture = reinterpret_cast<ID3D11Resource*&> (resource);

            ID3D11ShaderResourceView* shaderResourceView;
            const HRESULT result = CreateDDSTextureFromFile(&device, filenames[filenameIndex].c_str(), &texture, &shaderResourceView);
            DxErrorChecker(result);

            shaderResourceView->Release();
        }

        //
        // Create the texture array.  Each element in the texture 
        // array has the same format/dimensions.
        //

        D3D11_TEXTURE2D_DESC textureElementDesc;
        sourceTextures[0]->GetDesc(&textureElementDesc);

        D3D11_TEXTURE2D_DESC textureArrayDesc;
        textureArrayDesc.Width = textureElementDesc.Width;
        textureArrayDesc.Height = textureElementDesc.Height;
        textureArrayDesc.MipLevels = textureElementDesc.MipLevels;
        textureArrayDesc.ArraySize = numFilenames;
        textureArrayDesc.Format = textureElementDesc.Format;
        textureArrayDesc.SampleDesc.Count = 1;
        textureArrayDesc.SampleDesc.Quality = 0;
        textureArrayDesc.Usage = D3D11_USAGE_DEFAULT;
        textureArrayDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
        textureArrayDesc.CPUAccessFlags = 0;
        textureArrayDesc.MiscFlags = 0;

        ID3D11Texture2D* textureArray;
        HRESULT result = device.CreateTexture2D(&textureArrayDesc, 0, &textureArray);
        DxErrorChecker(result);

        D3D11_TEXTURE2D_DESC textureElementDesc2;
        textureArray->GetDesc(&textureElementDesc2);
        

        //
        // Copy individual texture elements into texture array.
        //

        // for each texture element...
        const size_t mipLevels = textureElementDesc.MipLevels; 
        for (uint32_t filenameIndex = 0; filenameIndex < numFilenames; ++filenameIndex) {
            // for each mipmap level...
            for(uint32_t mipLevelIndex = 0; mipLevelIndex < mipLevels; ++mipLevelIndex) {
                const uint32_t subResourceIndex = D3D11CalcSubresource(mipLevelIndex, 0, static_cast<UINT> (mipLevels));
				const uint32_t destinationSubresource = D3D11CalcSubresource(mipLevelIndex, filenameIndex, static_cast<UINT> (mipLevels));
                context.CopySubresourceRegion(textureArray, static_cast<uint32_t> (destinationSubresource), 0, 0, 0, sourceTextures[filenameIndex], subResourceIndex, nullptr);
            }
        }	

        // Create a resource view to the texture array.
        D3D11_SHADER_RESOURCE_VIEW_DESC viewDesc;
        viewDesc.Format = textureArrayDesc.Format;
        viewDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2DARRAY;
        viewDesc.Texture2DArray.MostDetailedMip = 0;
        viewDesc.Texture2DArray.MipLevels = textureArrayDesc.MipLevels;
        viewDesc.Texture2DArray.FirstArraySlice = 0;
        viewDesc.Texture2DArray.ArraySize = numFilenames;

        result = device.CreateShaderResourceView(textureArray, &viewDesc, &resourceView);
        DxErrorChecker(result);

        // Cleanup--we only need the resource view.
        textureArray->Release();

        for(size_t filenameIndex = 0; filenameIndex < numFilenames; ++filenameIndex) {
            sourceTextures[filenameIndex]->Release();
        }
    }
}


ShaderResources::~ShaderResources() {
	for each (auto& elem in mShaderResourceViews) {
		elem.second->Release();
	}

	for each (auto& elem in mUnorderedAccessViews) {
		elem.second->Release();
	}

	for each (auto& elem in mBuffers) {
		elem.second->Release();
	}
}

ID3D11ShaderResourceView* 
ShaderResources::addTextureFromFileSRV(const std::wstring& name, const std::wstring& filepath) {
	// Names must be unique
	if (mShaderResourceViews.find(name) != mShaderResourceViews.end()) {
		return nullptr;
	}

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;

	ID3D11Resource* texture;
	ID3D11ShaderResourceView* srv;
	const HRESULT result = CreateDDSTextureFromFile(&device, filepath.c_str(), &texture, &srv);
	DxErrorChecker(result);

	texture->Release();

	mShaderResourceViews[name] = srv;

	return srv;
}

ID3D11ShaderResourceView* 
ShaderResources::addResourceSRV(const std::wstring& name, ID3D11Resource* resource, const D3D11_SHADER_RESOURCE_VIEW_DESC& desc) {
	// Names must be unique
	if (mShaderResourceViews.find(name) != mShaderResourceViews.end()) {
		return nullptr;
	}

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;

	ID3D11ShaderResourceView* view;
	const HRESULT result = device.CreateShaderResourceView(resource, &desc, &view);
	DxErrorChecker(result);
	
	mShaderResourceViews[name] = view;

	return view;
}

ID3D11UnorderedAccessView* 
ShaderResources::addResourceUAV(const std::wstring& name, ID3D11Resource* resource, const D3D11_UNORDERED_ACCESS_VIEW_DESC& desc) {
	// Names must be unique
	if (mUnorderedAccessViews.find(name) != mUnorderedAccessViews.end()) {
		return nullptr;
	}

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;

	ID3D11UnorderedAccessView* view;
	const HRESULT result = device.CreateUnorderedAccessView(resource, &desc, &view);
	DxErrorChecker(result);

	mUnorderedAccessViews[name] = view;

	return view;
}

ID3D11Buffer* 
ShaderResources::addBuffer(const std::wstring& name, D3D11_BUFFER_DESC& desc, const D3D11_SUBRESOURCE_DATA* initData) {
	// Names must be unique
	if (mBuffers.find(name) != mBuffers.end()) {
		return nullptr;
	}

	assert(Globals::gD3D->mDevice);
	ID3D11Device& device = *Globals::gD3D->mDevice;

	// If the bind flag is D3D11_BIND_CONSTANT_BUFFER, 
	// you must set the ByteWidth value in multiples of 16, 
	// and less than or equal to D3D11_REQ_CONSTANT_BUFFER_ELEMENT_COUNT.
	if (desc.ByteWidth == D3D11_BIND_CONSTANT_BUFFER) {
		const size_t baseAlignment = 16;
		desc.ByteWidth = static_cast<uint32_t>(desc.ByteWidth + (baseAlignment - (desc.ByteWidth % baseAlignment)));
	}

	ID3D11Buffer* buffer;
	const HRESULT result = device.CreateBuffer(&desc, initData, &buffer);
	DxErrorChecker(result);

	mBuffers[name] = buffer;

	return buffer;
}

ID3D11ShaderResourceView* 
ShaderResources::getShaderResourceView(const std::wstring& name) const {
	ShaderResourceViews::const_iterator it = mShaderResourceViews.find(name);
	if (it != mShaderResourceViews.end()) {
		return it->second;
	} else {
		return nullptr;
	}
}

ID3D11UnorderedAccessView* 
ShaderResources::getUnorderedAccessView(const std::wstring& name) const {
	UnorderedAccessViews::const_iterator it = mUnorderedAccessViews.find(name);
	if (it != mUnorderedAccessViews.end()) {
		return it->second;
	}
	else {
		return nullptr;
	}
}

ID3D11Buffer* 
ShaderResources::getBuffer(const std::wstring& name) const {
	Buffers::const_iterator it = mBuffers.find(name);
	if (it != mBuffers.end()) {
		return it->second;
	}
	else {
		return nullptr;
	}
}


